#include "robot.h"

#define EPSINON 1E-4
/*#####################################  joint limits  ###############################################*/
const static JointsLimits joint_limits = {
                            .joint1_limit_min = -PI,
                            .joint1_limit_max = PI,
                            .joint2_limit_min = -PI,
                            .joint2_limit_max = PI,
                            .joint3_limit_min = -PI,
                            .joint3_limit_max = PI,
                            .joint4_limit_min = -PI,
                            .joint4_limit_max = PI,
                            .joint5_limit_min = -PI,
                            .joint5_limit_max = PI,
                            .joint6_limit_min = -PI,
                            .joint6_limit_max = PI
                            };
/*#####################################  roboticArm model  ###############################################*/
static Robot roboticArm;
static int roboticArm_ready = 0;// a flag to check roboticArm model finish init.
void robot_init(){
    DH dh1 = {.alpha=0.0,           .a=0.0,         .init_theta=0.0,                 .d=0.0};
    DH dh2 = {.alpha=-PI/2,         .a=109,         .init_theta=0.0,                 .d=0.0};
    DH dh3 = {.alpha=0.0,           .a=250.5,       .init_theta=0.0,                 .d=0.0};
    DH dh4 = {.alpha=-PI/2,         .a=115.5,       .init_theta=0.0,                 .d=239.5};
    DH dh5 = {.alpha=PI/2,          .a=0.0,         .init_theta=0.0,                 .d=0.0};
    DH dh6 = {.alpha=-PI/2,         .a=0.0,         .init_theta=0.0,                 .d=0.0};
    DH dh_list[6] = {dh1, dh2, dh3, dh4, dh5, dh6};
    DH dh_base = {.alpha=0.0, .a=0.0, .init_theta=0.0, .d=0};
    DH dh_tcp = {.alpha=0.0, .a=0.0, .init_theta=0.0, .d=0};
    setRobotModel(dh_list, 6);
    setRobotBase(&dh_base);
    setRobotTcp(&dh_tcp);
    roboticArm_ready = 1;
}

void setRobotModel(DH* dh_list, int dof){
    roboticArm.dof = dof;
    roboticArm.dh_0_1 = dh_list[0];
    roboticArm.dh_1_2 = dh_list[1];
    roboticArm.dh_2_3 = dh_list[2];
    roboticArm.dh_3_4 = dh_list[3];
    roboticArm.dh_4_5 = dh_list[4];
    roboticArm.dh_5_6 = dh_list[5];
}

void setRobotBase(DH* dh){
    roboticArm.dh_base_0 = *dh;
}

void setRobotTcp(DH* dh){
    roboticArm.dh_end_tcp = *dh;
}

void getRobotTfTree(Joints* joints, Matrix* tf_tree){
    DH2Matrix(roboticArm.dh_0_1, joints->theta1, tf_tree);
    DH2Matrix(roboticArm.dh_1_2, joints->theta2, tf_tree+1);
    DH2Matrix(roboticArm.dh_2_3, joints->theta3, tf_tree+2);
    DH2Matrix(roboticArm.dh_3_4, joints->theta4, tf_tree+3);
    DH2Matrix(roboticArm.dh_4_5, joints->theta5, tf_tree+4);
    DH2Matrix(roboticArm.dh_5_6, joints->theta6, tf_tree+5);
}

void forward_kinematics(Joints* joints, Pose* pose){
    Matrix robot_tf_tree[6];
    getRobotTfTree(joints, robot_tf_tree);
    Matrix tf_0_1 = robot_tf_tree[0];
    Matrix tf_0_2, tf_0_3, tf_0_4, tf_0_5, tf_0_6;
    tf_0_2 = matrixMultiply(&tf_0_1, &robot_tf_tree[1]);
    tf_0_3 = matrixMultiply(&tf_0_2, &robot_tf_tree[2]);
    tf_0_4 = matrixMultiply(&tf_0_3, &robot_tf_tree[3]);
    tf_0_5 = matrixMultiply(&tf_0_4, &robot_tf_tree[4]);
    tf_0_6 = matrixMultiply(&tf_0_5, &robot_tf_tree[5]);
    matrix2Pose(&tf_0_6, pose);
}

void inverse_kinematics(Pose* target_pose, PossibleIK* ik_joints){
    Matrix target_matrix;
    pose2Matrix(target_pose, &target_matrix);
    float nx = target_matrix.m[0][0]; float ox = target_matrix.m[0][1]; float ax = target_matrix.m[0][2];
    float ny = target_matrix.m[1][0]; float oy = target_matrix.m[1][1]; float ay = target_matrix.m[1][2];
    float nz = target_matrix.m[2][0]; float oz = target_matrix.m[2][1]; float az = target_matrix.m[2][2];
    float px = target_matrix.m[0][3]; float py = target_matrix.m[1][3]; float pz = target_matrix.m[2][3];
    float a0 = roboticArm.dh_0_1.a; float d1 = roboticArm.dh_0_1.d; 
    float a1 = roboticArm.dh_1_2.a; float d2 = roboticArm.dh_1_2.d; 
    float a2 = roboticArm.dh_2_3.a; float d3 = roboticArm.dh_2_3.d; 
    float a3 = roboticArm.dh_3_4.a; float d4 = roboticArm.dh_3_4.d; 
    float a4 = roboticArm.dh_4_5.a; float d5 = roboticArm.dh_4_5.d; 
    float a5 = roboticArm.dh_5_6.a; float d6 = roboticArm.dh_5_6.d; 
    /************************ joint 1 *****************************/
    float possible_theta1[2] = {0, 0};
    possible_theta1[0] = -atan2(-py, px);
    possible_theta1[1] = -atan2(-py, px) + PI;
    /************************ joint 3 *****************************/
    float possible_theta3[4] = {0, 0, 0, 0};
    /* temp, x, y, z, A are temp valiables */
    float temp[2] = {0, 0};
    float A[2] = {0, 0};
    float x = a2 * d4;
    float y = a2 * a3;
    float z = a2 * a2 + a3 * a3 + d4 * d4;
    int i;
    for (i = 0; i < 4; i+=2) {
        temp[i/2] = a1 - px * cos(possible_theta1[i/2]) - py * sin(possible_theta1[i/2]);
        A[i/2] = z - pz * pz - temp[i/2] * temp[i/2];
        possible_theta3[i]   = -atan2(-y, x) + atan2(A[i/2], sqrt( (4 * x * x + 4 * y * y - A[i/2] * A[i/2])));
        possible_theta3[i+1] = -atan2(-y, x) + atan2(A[i/2], -sqrt((4 * x * x + 4 * y * y - A[i/2] * A[i/2])));
    }
    /************************ joint 2 *****************************/
    float possible_theta2[4] = {0, 0, 0, 0};
    float m[4] = {0, 0, 0, 0};
    float n[4] = {0, 0, 0, 0};
    for (i = 0; i < 4; ++i) {
        m[i] = a2 + a3 * cos(possible_theta3[i]) - d4 * sin(possible_theta3[i]);
        n[i] = a3 * sin(possible_theta3[i]) + d4 * cos(possible_theta3[i]);
        possible_theta2[i] = atan2(-pz * m[i] + n[i] * temp[i/2],-pz * n[i] - m[i] * temp[i/2]);
    }
    /************************ joint 5 *****************************/
    float possible_theta5[8] = {0, 0, 0, 0, 0, 0, 0, 0};
    for (i = 0; i < 4; ++i) {
        m[i] = (ax * cos(possible_theta1[i / 2]) * cos(possible_theta2[i]) + ay * sin(possible_theta1[i / 2]) * cos(possible_theta2[i]) + az * -1 * sin(possible_theta2[i]));
        n[i] = -(ax * cos(possible_theta1[i / 2]) * sin(possible_theta2[i]) + ay * sin(possible_theta1[i / 2]) * sin(possible_theta2[i]) - az * -1 * cos(possible_theta2[i]));
    }
    for (i = 0; i < 8; ) {
        x = n[i/2] * cos(possible_theta3[i/2]) - m[i/2] * sin(possible_theta3[i/2]);
        y = sqrt(pow((ay * cos(possible_theta1[i / 4]) - ax * sin(possible_theta1[i / 4])),2) + pow((m[i/2] * cos(possible_theta3[i/2]) + n[i/2] * sin(possible_theta3[i/2])),2));
        possible_theta5[i++] = atan2( y, x);
        possible_theta5[i++] = atan2(-y, x);
    }
    /************************ joint 4 *****************************/
    float possible_theta4[8] = {0, 0, 0, 0, 0, 0, 0, 0};
    for (i = 0; i < 8; ++i) {
        if (fabs(sin(possible_theta5[i])) >= EPSINON) {
            y = ((ay*cos(possible_theta1[i/4])-ax*sin(possible_theta1[i/4]))) / sin(possible_theta5[i]);
            x= (-m[i/2]*cos(possible_theta3[i/2])-n[i/2]*sin(possible_theta3[i/2]))/(sin(possible_theta5[i]));
            possible_theta4[i] = atan2(y , x);
        } 
        else {
            possible_theta4[i] = 0;
        }
    }
    /************************ joint 6 *****************************/
    float possible_theta6[8] = {0, 0, 0, 0, 0, 0, 0 ,0};
    for (i = 0; i < 8; ++i) {
        float u = nx * sin(possible_theta1[i/4]) - ny * cos(possible_theta1[i/4]);
        float v = ox * sin(possible_theta1[i/4]) - oy * cos(possible_theta1[i/4]);
        possible_theta6[i] = atan2(cos(possible_theta4[i]) * u - cos(possible_theta5[i]) * sin(possible_theta4[i]) * v, cos(possible_theta4[i]) * v + cos(possible_theta5[i]) * sin(possible_theta4[i]) * u);
    }
    /*************************assign to 8 possible IK joints*********************************/
    // 0 0 0 0 0 0 
    ik_joints->ik_results[0].theta1 = possible_theta1[0];
    ik_joints->ik_results[0].theta2 = possible_theta2[0];
    ik_joints->ik_results[0].theta3 = possible_theta3[0];
    ik_joints->ik_results[0].theta4 = possible_theta4[0];
    ik_joints->ik_results[0].theta5 = possible_theta5[0];
    ik_joints->ik_results[0].theta6 = possible_theta6[0];
    // 0 0 0 1 1 1
    ik_joints->ik_results[1].theta1 = possible_theta1[0];
    ik_joints->ik_results[1].theta2 = possible_theta2[0];
    ik_joints->ik_results[1].theta3 = possible_theta3[0];
    ik_joints->ik_results[1].theta4 = possible_theta4[1];
    ik_joints->ik_results[1].theta5 = possible_theta5[1];
    ik_joints->ik_results[1].theta6 = possible_theta6[1];
    // 0 1 1 2 2 2
    ik_joints->ik_results[2].theta1 = possible_theta1[0];
    ik_joints->ik_results[2].theta2 = possible_theta2[1];
    ik_joints->ik_results[2].theta3 = possible_theta3[1];
    ik_joints->ik_results[2].theta4 = possible_theta4[2];
    ik_joints->ik_results[2].theta5 = possible_theta5[2];
    ik_joints->ik_results[2].theta6 = possible_theta6[2];
    // 0 1 1 3 3 3
    ik_joints->ik_results[3].theta1 = possible_theta1[1];
    ik_joints->ik_results[3].theta2 = possible_theta2[1];
    ik_joints->ik_results[3].theta3 = possible_theta3[1];
    ik_joints->ik_results[3].theta4 = possible_theta4[3];
    ik_joints->ik_results[3].theta5 = possible_theta5[3];
    ik_joints->ik_results[3].theta6 = possible_theta6[3];
    // 1 2 2 4 4 4
    ik_joints->ik_results[4].theta1 = possible_theta1[1];
    ik_joints->ik_results[4].theta2 = possible_theta2[2];
    ik_joints->ik_results[4].theta3 = possible_theta3[2];
    ik_joints->ik_results[4].theta4 = possible_theta4[4];
    ik_joints->ik_results[4].theta5 = possible_theta5[4];
    ik_joints->ik_results[4].theta6 = possible_theta6[4];
    // 1 2 2 5 5 5
    ik_joints->ik_results[5].theta1 = possible_theta1[1];
    ik_joints->ik_results[5].theta2 = possible_theta2[2];
    ik_joints->ik_results[5].theta3 = possible_theta3[2];
    ik_joints->ik_results[5].theta4 = possible_theta4[5];
    ik_joints->ik_results[5].theta5 = possible_theta5[5];
    ik_joints->ik_results[5].theta6 = possible_theta6[5];
    // 1 3 3 6 6 6 
    ik_joints->ik_results[6].theta1 = possible_theta1[1];
    ik_joints->ik_results[6].theta2 = possible_theta2[3];
    ik_joints->ik_results[6].theta3 = possible_theta3[3];
    ik_joints->ik_results[6].theta4 = possible_theta4[6];
    ik_joints->ik_results[6].theta5 = possible_theta5[6];
    ik_joints->ik_results[6].theta6 = possible_theta6[6];
    // 1 3 3 7 7 7
    ik_joints->ik_results[7].theta1 = possible_theta1[1];
    ik_joints->ik_results[7].theta2 = possible_theta2[3];
    ik_joints->ik_results[7].theta3 = possible_theta3[3];
    ik_joints->ik_results[7].theta4 = possible_theta4[7];
    ik_joints->ik_results[7].theta5 = possible_theta5[7];
    ik_joints->ik_results[7].theta6 = possible_theta6[7];
}

static float find_max_joint_diff(Joints* joint1, Joints* joint2){
    float max_diff = -1000.0;
    float diff;
    diff = fabs(joint1->theta1 - joint2->theta1);
    max_diff = diff > max_diff ? diff : max_diff;
    diff = fabs(joint1->theta2 - joint2->theta2);
    max_diff = diff > max_diff ? diff : max_diff;
    diff = fabs(joint1->theta3 - joint2->theta3);
    max_diff = diff > max_diff ? diff : max_diff;
    diff = fabs(joint1->theta4 - joint2->theta4);
    max_diff = diff > max_diff ? diff : max_diff;
    diff = fabs(joint1->theta5 - joint2->theta5);
    max_diff = diff > max_diff ? diff : max_diff;
    diff = fabs(joint1->theta6 - joint2->theta6);
    max_diff = diff > max_diff ? diff : max_diff;
    return max_diff;
}

int compute_best_ik(PossibleIK* ik_joints, Joints* start_joints, Joints* best_ik){
    int result = 0;
    int flag[8] = {0};
    int i;
    for(i = 0; i < 8; i++){
        // check joint isnan
        if(isnan(ik_joints->ik_results[i].theta1)  || isnan(ik_joints->ik_results[i].theta2) ||
            isnan(ik_joints->ik_results[i].theta3) || isnan(ik_joints->ik_results[i].theta4) ||
            isnan(ik_joints->ik_results[i].theta5) || isnan(ik_joints->ik_results[i].theta6) )  {
                continue;
            }
        // check joint whether whthin joint limit
        if( (ik_joints->ik_results[i].theta1 >= joint_limits.joint1_limit_min) && 
            (ik_joints->ik_results[i].theta1 <= joint_limits.joint1_limit_max) &&
            (ik_joints->ik_results[i].theta2 >= joint_limits.joint2_limit_min) &&
            (ik_joints->ik_results[i].theta2 <= joint_limits.joint2_limit_max) &&
            (ik_joints->ik_results[i].theta3 >= joint_limits.joint3_limit_min) &&
            (ik_joints->ik_results[i].theta3 <= joint_limits.joint3_limit_max) &&
            (ik_joints->ik_results[i].theta4 >= joint_limits.joint4_limit_min) &&
            (ik_joints->ik_results[i].theta4 <= joint_limits.joint4_limit_max) &&
            (ik_joints->ik_results[i].theta5 >= joint_limits.joint5_limit_min) &&
            (ik_joints->ik_results[i].theta5 <= joint_limits.joint5_limit_max) &&
            (ik_joints->ik_results[i].theta6 >= joint_limits.joint6_limit_min) &&
            (ik_joints->ik_results[i].theta6 <= joint_limits.joint6_limit_max)){
                flag[i] = 1;
                result = 1;
            }
    }
    if(result){
        int best_ik_index = -1;
        float min_max_joint_angle = 1000.0;
        for(i = 0; i < 8; i++){
            if(flag[i] == 0) continue;
            float temp_max = find_max_joint_diff(&(ik_joints->ik_results[i]), start_joints);
            if(temp_max < min_max_joint_angle){
                min_max_joint_angle = temp_max;
                best_ik_index = i;
            }
        best_ik->theta1 = ik_joints->ik_results[best_ik_index].theta1;
        best_ik->theta2 = ik_joints->ik_results[best_ik_index].theta2;
        best_ik->theta3 = ik_joints->ik_results[best_ik_index].theta3;
        best_ik->theta4 = ik_joints->ik_results[best_ik_index].theta4;
        best_ik->theta5 = ik_joints->ik_results[best_ik_index].theta5;
        best_ik->theta6 = ik_joints->ik_results[best_ik_index].theta6;
        }
    }
    return result;
}

void get_best_ik(Pose* target_pose, Joints* start_joints, Joints* best_ik){
    PossibleIK possible_joints;
    inverse_kinematics(target_pose, &possible_joints);
    compute_best_ik(&possible_joints, start_joints, best_ik);
    return;
}

int jointSpacePlanner(Joints* start_joints, Joints* end_joints, Joints* mid_joints_list, int max_num){
    static float step1 = 5 * PI/ 180; //interpolation accuracy: 5deg
    static float step2 = 5 * PI/ 180; //interpolation accuracy: 5deg
    static float step3 = 4 * PI/ 180; //interpolation accuracy: 4deg
    static float step4 = 4 * PI/ 180; //interpolation accuracy: 4deg
    static float step5 = 3 * PI/ 180; //interpolation accuracy: 3deg
    static float step6 = 3 * PI/ 180; //interpolation accuracy: 3deg
    int joints_num = 0;
    int cnt_list[6];
    cnt_list[0] = fabs(end_joints->theta1 - start_joints->theta1) / step1 + 1;
    cnt_list[1] = fabs(end_joints->theta2 - start_joints->theta2) / step2 + 1;
    cnt_list[2] = fabs(end_joints->theta3 - start_joints->theta3) / step3 + 1;
    cnt_list[3] = fabs(end_joints->theta4 - start_joints->theta4) / step4 + 1;
    cnt_list[4] = fabs(end_joints->theta5 - start_joints->theta5) / step5 + 1;
    cnt_list[5] = fabs(end_joints->theta6 - start_joints->theta6) / step6 + 1;
    int max_cnt = -1;
    int i;
    for(i = 0; i < 6; i++){
        if(cnt_list[i] > max_cnt) {
            max_cnt = cnt_list[i];
        }
    }
    if(max_cnt > max_num){
        max_cnt = max_num;
    }
    float t;
    for(i = 0; i <= max_cnt; i++){
        t = (1.0 * i) / max_cnt;
        mid_joints_list[i].theta1 = (1-t)*(start_joints->theta1) + t*(end_joints->theta1);
        mid_joints_list[i].theta2 = (1-t)*(start_joints->theta2) + t*(end_joints->theta2);
        mid_joints_list[i].theta3 = (1-t)*(start_joints->theta3) + t*(end_joints->theta3);
        mid_joints_list[i].theta4 = (1-t)*(start_joints->theta4) + t*(end_joints->theta4);
        mid_joints_list[i].theta5 = (1-t)*(start_joints->theta5) + t*(end_joints->theta5);
        mid_joints_list[i].theta6 = (1-t)*(start_joints->theta6) + t*(end_joints->theta6);
        joints_num++;
    }
    return joints_num;
}